from sqlalchemy import Column, String, Integer
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from init.mysql import engine

Base = declarative_base()


class CompanyEmployee(Base):
    __tablename__ = 'company_employee'

    code = Column(String(6), primary_key=True)
    name = Column(String(20))
    title = Column(String(20))
    sdate = Column(String(10))
    edate = Column(String(10))
    type = Column(Integer)

    def __init__(self, code, name, title, sdate, edate, type):
        self.code = code
        self.name = name
        self.title = title
        self.sdate = sdate
        self.edate = edate
        self.type = type

    @staticmethod
    def insert_or_update(code, data_list, type):
        session = sessionmaker(bind=engine)()
        for data in data_list:
            existing_employee = session.query(CompanyEmployee).filter_by(code=code).first()
            if existing_employee:
                existing_employee.name = data['name']
                existing_employee.title = data['title']
                existing_employee.sdate = data['sdate']
                existing_employee.edate = data['edate']
                existing_employee.type = type
            else:
                session.add(CompanyEmployee(code, data['name'], data['title'], data['sdate'], data['edate'], type))
        session.commit()


Base.metadata.create_all(engine)
